import os
from dotenv import load_dotenv, find_dotenv # 导入 find_dotenv 帮助定位
from langchain.memory import ConversationSummaryBufferMemory
from langchain.llms import OpenAI

# 加载 .env 文件中的环境变量 (增强调试)
load_dotenv(dotenv_path=find_dotenv(usecwd=True), verbose=True, override=True)

# 从环境变量加载 API 密钥和基础 URL
api_key = os.getenv("OPENAI_API_KEY")
api_base = os.getenv("OPENAI_API_BASE")
os.environ["OPENAI_API_KEY"] = api_key
os.environ["OPENAI_API_BASE"] = api_base

# 当对话持续进行且对话内容很多的时候
# 可以使用ConversationSummaryBufferMemory来存储对话摘要
# 这是一种非常有用的方式,它会根据token的数量来自动判断是否需要进行摘要
# 当token数量超过阈值的时候,会自动进行摘要
# 在缓冲区中,会保留最近的k条对话
#比较久的对话会被删除，在删除前会进行摘要

memory = ConversationSummaryBufferMemory(
    #k=2,
    llm=OpenAI(temperature=0),
    max_token_limit=100,
    return_messages=True
)
memory.save_context(
    {"input":"帮我找一下tomie"},
    {"output":"对不起请问什么是tomie？"}
)
memory.save_context(
    {"input":"tomie是一个培训讲师"},
    {"output":"好的，我知道了。"}
)
memory.save_context(
    {"input":"今天他要讲一门关于RAG的课程"},
    {"output":"好的，我知道了。需要RAG的资料吗？"}
)

print(memory.load_memory_variables({}))